/**
 * @file IsReal_custom.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "is_real_custom_tiling.h"
#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

constexpr uint32_t IsReal_BFLOAT16 = 0;            // 对应 DT_BF16，bf16 类型
constexpr uint32_t IsReal_FLOAT16 = 1;             // 对应 DT_FLOAT16，half（float16）类型
constexpr uint32_t IsReal_FLOAT32 = 2;             // 对应 DT_FLOAT，float32 类型
constexpr uint32_t IsReal_INT8 = 3;                // 对应 DT_INT8，int8 类型
constexpr uint32_t IsReal_INT16 = 4;               // 对应 DT_INT16，int16 类型
constexpr uint32_t IsReal_INT32 = 5;               // 对应 DT_INT32，int32 类型
constexpr uint32_t IsReal_UINT8 = 6;               // 对应 DT_UINT8，u8（无符号int8）类型
constexpr uint32_t IsReal_UINT16 = 7;              // 对应 DT_UINT16，u16（无符号int16）类型
constexpr uint32_t IsReal_UINT32 = 8;              // 对应 DT_UINT32，u32（无符号int32）类型
constexpr uint32_t IsReal_INT64 = 9;               // 对应 DT_INT64，int64 类型
constexpr uint32_t IsReal_UINT64 = 10;             // 对应 DT_UINT64，u64（无符号int64）类型
constexpr uint32_t IsReal_BOOL = 11;               // 对应 DT_BOOL，bool 类型
constexpr uint32_t IsReal_DOUBLE = 12;             // 对应 DT_DOUBLE，double（float64）类型
constexpr uint32_t IsReal_STRING = 13;             // 对应 DT_STRING，string（字符串）类型
constexpr uint32_t IsReal_DUAL_SUB_INT8 = 14;      // 对应 DT_DUAL_SUB_INT8，dual output int8（双输出int8）类型
constexpr uint32_t IsReal_DUAL_SUB_UINT8 = 15;     // 对应 DT_DUAL_SUB_UINT8，dual output u8（双输出无符号int8）类型
constexpr uint32_t IsReal_COMPLEX64 = 16;          // 对应 DT_COMPLEX64，complex64（复数64位）类型
constexpr uint32_t IsReal_COMPLEX128 = 17;         // 对应 DT_COMPLEX128，complex128（复数128位）类型
constexpr uint32_t IsReal_QINT8 = 18;              // 对应 DT_QINT8，qint8（量化int8）类型
constexpr uint32_t IsReal_QINT16 = 19;             // 对应 DT_QINT16，qint16（量化int16）类型
constexpr uint32_t IsReal_QINT32 = 20;             // 对应 DT_QINT32，qint32（量化int32）类型
constexpr uint32_t IsReal_QUINT8 = 21;             // 对应 DT_QUINT8，quint8（量化无符号int8）类型
constexpr uint32_t IsReal_QUINT16 = 22;            // 对应 DT_QUINT16，quint16（量化无符号int16）类型
constexpr uint32_t IsReal_RESOURCE = 23;           // 对应 DT_RESOURCE，resource（资源）类型
constexpr uint32_t IsReal_STRING_REF = 24;         // 对应 DT_STRING_REF，string ref（字符串引用）类型
constexpr uint32_t IsReal_DUAL = 25;               // 对应 DT_DUAL，dual output（双输出）类型
constexpr uint32_t IsReal_VARIANT = 26;            // 对应 DT_VARIANT，dt_variant（变体）类型
constexpr uint32_t IsReal_UNDEFINED = 27;          // 对应 DT_UNDEFINED，未设置的DataType字段
constexpr uint32_t IsReal_INT4 = 28;               // 对应 DT_INT4，int4 类型
constexpr uint32_t IsReal_UINT1 = 29;              // 对应 DT_UINT1，u1（无符号int1）类型
constexpr uint32_t IsReal_INT2 = 30;               // 对应 DT_INT2，int2 类型
constexpr uint32_t IsReal_UINT2 = 31;              // 对应 DT_UINT2，u2（无符号int2）类型
constexpr uint32_t IsReal_COMPLEX32 = 32;          // 对应 DT_COMPLEX32，complex32（复数32位）类型
constexpr uint32_t IsReal_HIFLOAT8 = 33;           // 对应 DT_HIFLOAT8，hif8（高精度float8）类型
constexpr uint32_t IsReal_FLOAT8_E5M2 = 34;        // 对应 DT_FLOAT8_E5M2，fp8_e5m2（float8格式，5位指数2位尾数）类型
constexpr uint32_t IsReal_FLOAT8_E4M3FN = 35;      // 对应 DT_FLOAT8_E4M3FN，fp8_e4m3（float8格式，4位指数3位尾数）类型
constexpr uint32_t IsReal_FLOAT8_E8M0 = 36;        // 对应 DT_FLOAT8_E8M0，fp8_e8m0（float8格式，8位指数0位尾数）类型
constexpr uint32_t IsReal_FLOAT6_E3M2 = 37;        // 对应 DT_FLOAT6_E3M2，fp6_e3m2（float6格式，3位指数2位尾数）类型
constexpr uint32_t IsReal_FLOAT6_E2M3 = 38;        // 对应 DT_FLOAT6_E2M3，fp6_e2m3（float6格式，2位指数3位尾数）类型
constexpr uint32_t IsReal_FLOAT4_E2M1 = 39;        // 对应 DT_FLOAT4_E2M1，fp4_e2m1（float4格式，2位指数1位尾数）类型
constexpr uint32_t IsReal_FLOAT4_E1M2 = 40;        // 对应 DT_FLOAT4_E1M2，fp4_e1m2（float4格式，1位指数2位尾数）类型
constexpr uint32_t IsReal_MAX = 41;                // 对应 DT_MAX，标记AscendC数据类型的边界

constexpr uint32_t LAST_TWO_TILE = 2;

// 定义复数类型
struct complex64 {
    float real;
    float imag;
};

// 定义复数类型
struct complex128 {
    double real;
    double imag;
};

template <typename dataType> class KernelIsReal {
public:
    __aicore__ inline KernelIsReal() {}
    __aicore__ inline void Init(GM_ADDR x,GM_ADDR z, IsRealCustomTilingData tiling)
    {
        this->dataTypeId = tiling.dataType;
        if (tiling.isEvenCore) {
            this->blockLength = tiling.blockLength;
            this->tileNum = tiling.tileNum;
            this->tileLength = tiling.tileLength / BUFFER_NUM;
            this->lastTileLength = tiling.lastTileLength;
            xGm.SetGlobalBuffer((__gm__ dataType *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
            zGm.SetGlobalBuffer((__gm__ int8_t *)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        } else {
            if (AscendC::GetBlockIdx() < tiling.formerNum) {
                this->tileNum = tiling.formerTileNum;
                this->tileLength = tiling.formerTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.formerLastTileLength;
                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * AscendC::GetBlockIdx(), tiling.formerLength);
                zGm.SetGlobalBuffer((__gm__ int8_t *)z + tiling.formerLength * AscendC::GetBlockIdx(), tiling.formerLength);
            } else {
                this->tileNum = tiling.tailTileNum;
                this->tileLength = tiling.tailTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.tailLastTileLength;

                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * tiling.formerNum + 
                    tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum), tiling.tailLength);
                zGm.SetGlobalBuffer((__gm__ int8_t *)z + tiling.formerLength * tiling.formerNum + 
                    tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum), tiling.tailLength);
            }
        }
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(int8_t));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<dataType> xLocal = inQueueX.AllocTensor<dataType>();
        if (progress == (this->tileNum * BUFFER_NUM - 1)) {
            AscendC::DataCopy(xLocal, xGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength],
                this->tileLength);
        } else {
            AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        }
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<dataType> xLocal = inQueueX.DeQue<dataType>();
        AscendC::LocalTensor<int8_t> zLocal = outQueueZ.AllocTensor<int8_t>();

        bool isComplex = (dataTypeId == IsReal_COMPLEX64 || dataTypeId == IsReal_COMPLEX128);
        
        for (int i = 0; i < this->tileLength; ++i) {
            zLocal.SetValue(i, isComplex ? 0 : 1);
            // 临时调试：强制输出0
            // zLocal.SetValue(i,  0 );
        }

        outQueueZ.EnQue<int8_t>(zLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<int8_t> zLocal = outQueueZ.DeQue<int8_t>();
        if (progress == (this->tileNum * BUFFER_NUM - 1)) {
            AscendC::DataCopy(zGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength], zLocal,
                this->tileLength);
        } else {
            AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
        }
        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;

    AscendC::GlobalTensor<dataType> xGm;
    AscendC::GlobalTensor<int8_t> zGm;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;
    uint32_t dataTypeId;
};
extern "C" __global__ __aicore__ void is_real_custom(GM_ADDR x,GM_ADDR z, IsRealCustomTilingData tiling)
{
    if (tiling.dataType == IsReal_BFLOAT16) {
        KernelIsReal<bfloat16_t> op;
        op.Init(x, z, tiling);
        op.Process();
    } else if (tiling.dataType == IsReal_FLOAT16) {
        KernelIsReal<half> op;
        op.Init(x, z, tiling);
        op.Process();
    } else if (tiling.dataType == IsReal_FLOAT32) {
        KernelIsReal<float> op;
        op.Init(x, z, tiling);
        op.Process();
    } else if (tiling.dataType == IsReal_INT8) {
        KernelIsReal<int8_t> op;
        op.Init(x, z, tiling);
        op.Process();
    } else if (tiling.dataType == IsReal_INT16) {
        KernelIsReal<int16_t> op;
        op.Init(x, z, tiling);
        op.Process();
    } else if (tiling.dataType == IsReal_INT32) {
        KernelIsReal<int32_t> op;
        op.Init(x, z, tiling);
        op.Process();
    }else if (tiling.dataType == IsReal_COMPLEX64) {
        KernelIsReal<complex64> op;
        op.Init(x, z, tiling);
        op.Process();
    }else if (tiling.dataType == IsReal_COMPLEX128) {
        KernelIsReal<complex128> op;
        op.Init(x, z, tiling);
        op.Process();
    } 
    else {
        return;
    }
}
